# Minimal stub for NVIDIA Apex so projects importing `from apex import amp`
# can run in pure FP32 without real Apex.
class _AmpStub:
    def float_function(self, fn):
        # No-op decorator: return the function unchanged
        return fn

    def scale_loss(self, loss, optimizer):
        # No-op context manager: yields the original loss
        class _Ctx:
            def __enter__(self_inner): return loss
            def __exit__(self_inner, exc_type, exc, tb): return False
        return _Ctx()

    def initialize(self, model, optimizer=None, **kwargs):
        # Just return model/optimizer unchanged
        return model, optimizer

# export an `amp` object compatible with `from apex import amp`
amp = _AmpStub()

# OPTIONAL: also expose top-level `initialize` for `from apex import amp; amp.initialize(...)`
initialize = amp.initialize
